package com.demo.ai.controller;

import com.demo.ai.functions.WeatherFunction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.time.LocalDate;

@RestController
@RequestMapping("/ai")
public class AiController {

    private final ChatClient chatClient;

    private final VectorStore vectorStore;

    public AiController(ChatClient.Builder chatClientBuilder,
                        ChatMemory chatMemory,
                        ToolCallbackProvider tools,
                        VectorStore vectorStore) {
        this.vectorStore = vectorStore;
        this.chatClient = chatClientBuilder
                //系统提示词
                .defaultSystem("""
                        ##角色
                        你是一个Java架构师，请以友好的方式回复Java相关的问题。
                        今天的日期是 {current_date}
                        """)
                //chatMemory对话记忆
                .defaultAdvisors(
                        PromptChatMemoryAdvisor.builder(chatMemory).build(),
                        new SimpleLoggerAdvisor()
                )
                //FunctionCalling
//                .defaultTools(new WeatherFunction())
                //mcp-stdio
                .defaultToolCallbacks(tools)
                //默认的RAG增强提示
                .defaultAdvisors(QuestionAnswerAdvisor.builder(vectorStore)
                        .searchRequest(SearchRequest.builder().similarityThreshold(0.8d).topK(1).build())
                        .build())
                .build();
    }

    @GetMapping(value = "/url/get", produces = "text/html;charset=UTF-8")
    String urlGet(String question) {
        return this.chatClient.prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .user(question)
                .system(item -> item.param("current_date", LocalDate.now().toString()))
                .call()
                .content();
    }

    @GetMapping(value = "/url/get/stream", produces = "text/html;charset=UTF-8")
    public Flux<String> urlGetStream(@RequestParam(value = "question", defaultValue = "你好") String question) {
        return chatClient.prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .user(question)
                .system(item -> item.param("current_date", LocalDate.now().toString()))
                .stream()
                .content();
    }

    @GetMapping(value = "/content/get/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> contentGetStream(@RequestParam(value = "question", defaultValue = "你好") String question) {
        return chatClient.prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .user(question)
                .system(item -> item.param("current_date", LocalDate.now().toString()))
                .stream()
                .content();
    }

}
